import cv2
import os
from utils.visualize import vis, label2chinese
import time
from detector_yolov7 import Detector

def detect_cap(file, weight='./weights/2022-12-13-bridge.pt', save_picture_dir='',
               save_video_path='', save_fourcc_format='mp4v'):
    """
    Args:
        file: 视频文件
        weight: 模型加载权重
        save_picture_dir: 处理完成的图片保存地址
        save_video_path: 处理完成的视频保存地址
        save_fourcc_format: 视频存储格式
    """
    # 获取图片文件夹及文件全名
    file_dir, file_name = os.path.split(file)
    # 获取视频文件的文件，名称与类型名称
    file_prefix, file_suffix = os.path.splitext(file_name)
    # 判断是否需要保存图片，如果需要，新建指定文件夹的picture文件夹并保存
    if save_picture_dir != '':
        save_picture_dir = os.path.join(save_picture_dir, 'picture')
        if not os.path.exists(save_picture_dir):
            os.mkdir(save_picture_dir)
    # 读取视频文件
    cap = cv2.VideoCapture(file)
    print(cap.isOpened())
    print('FRAME_HEIGHT:', cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print('FRAME_WIDTH:', cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    cap_fps = round(cap.get(cv2.CAP_PROP_FPS))
    # 设置视频文件每秒帧数
    # cap_fps = 1
    print('PROP_FPS:', cap.get(cv2.CAP_PROP_FPS))
    # 获取原视频总帧数
    frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    # frame_count = 1
    print("all frame number: ", frame_count)
    print(save_video_path)
    # if save_video_path != '':
    # 设置视频文件写入格式
    fourcc = cv2.VideoWriter_fourcc(*save_fourcc_format)
    # 设置视频文件写入的每一帧图片的宽高
    cap_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    # 定义opencv中的视频写入操作
    video_writer = cv2.VideoWriter(save_video_path, fourcc, cap_fps, cap_size)
    # 设置视频中截取图片跳跃的帧数
    jump_fps = int(cap.get(cv2.CAP_PROP_FPS))
    # 实例化检测器
    detector = Detector(weight=weight,
                        conf_thres=0.5,
                        iou_thres=0.5,
                        device_id="0")
    frame_num = 0
    # previous_frame = -jump_fps
    t0 = time.time()
    # while frame_num < 30 * 8 * cap_fps:
    while True:
        _, im = cap.read()
        if im is None:
            break
        frame_num = frame_num + 1

        # if not frame_num % jump_fps == 0:
        #     continue
        print('current frame id:', frame_num)
        print('current total time:', time.time() - t0)
        print('current process:', 100 * frame_num / frame_count, '%')
        time_start = time.time()
        _, detect_info = detector.detect(im)
        time_stop = time.time()
        print("detect_info time:", time_stop - time_start, " ms")
        if detect_info['box_nums'] > 0:
            detect_info['visual'] = vis(detect_info['raw_img'], detect_info['boxes'],
                                        detect_info['scores'], detect_info['class_ids'],
                                        detector.names)
            cv2.putText(detect_info['visual'], 'Frame ID' + ':' + str(frame_num), (600, 100),
                        cv2.FONT_HERSHEY_COMPLEX, 3.0, (300, 0, 0), 2)
        else:
            detect_info['visual'] = detect_info['raw_img']
            cv2.putText(detect_info['visual'], 'Frame ID' + ':' + str(frame_num), (600, 100),
                        cv2.FONT_HERSHEY_COMPLEX, 3.0, (300, 0, 0), 2)
        if save_video_path != '':
            video_writer.write(detect_info['visual'])
            # video_writer.write(im)
            # cv2.imwrite(f"./data/256A_{frame_num}.jpg", detect_info['visual'])
            cv2.imshow('detect_info', detect_info['visual'])
            cv2.waitKey(1)
    cap.release()
    video_writer.release()


if __name__ == '__main__':
    # storage_path = "/home/stark/algo-env/datasets/sample-chock/"
    # video_list = os.listdir(storage_path)
    # for video in video_list:
    #     video_path = os.path.join(storage_path, video)
    #     print("./demo-" + video)
    #     detect_cap(video_path, weight='weights/2022-12-05-chock.pt',
    #                save_picture_dir="",
    #                save_video_path="./demo-" + video,
    #                save_fourcc_format='mp4v')
    ############################################################################
    video_path = "data/172-03-01-06-37.mp4"
    detect_cap(video_path, weight='./weights/yolov7-channel.pt',
               save_picture_dir="",
               save_video_path="./done_test.mp4",
               save_fourcc_format='mp4v')
